#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <signal.h>
#include <netdb.h>
#include <sys/errno.h>
#include <sys/time.h>
#include <netinet/ip_icmp.h>
#include <netinet/ip.h>
#include <errno.h>

#define MAX_PACKETS 4
#define PACKET_SIZE 4096
#define DATA_LEN 56

struct sockaddr_in fromAddr = {0};
int nsend = 0;
int nrecved = 0;
struct timeval recvTimeVal;
char sendPacket[PACKET_SIZE];
char recvPacket[PACKET_SIZE];

// 函数声明
void send_packet(int fd, struct sockaddr_in* pDestAddr, int len);
int pack(int nsend);
u_short ckcsum(u_short* p_icmp, int pack_size);
void recv_packet(int sfd);
int unpack(int len, uid_t uid);
void tv_sub(struct timeval* out, struct timeval* in);
void print_statistics();

int main(int argc, char* argv[]) {
    if (argc < 2) {
        printf("用法: %s <IP地址或域名>\n", argv[0]);
        exit(-1);
    }

    // 获取ICMP协议号
    struct protoent* protocol = getprotobyname("icmp");
    if (NULL == protocol) {
        perror("getprotobyname失败");
        exit(-1);
    }

    // 创建原始套接字（需要root权限）
    int sfd = socket(AF_INET, SOCK_RAW, protocol->p_proto);
    if (sfd == -1) {
        perror("创建socket失败");
        exit(-1);
    }

    // 设置套接字选项
    int size = 1024 * 50;
    if (setsockopt(sfd, SOL_SOCKET, SO_RCVBUF, &size, sizeof(size)) == -1) {
        perror("设置socket失败");
        close(sfd);
        exit(-1);
    }

    // 解析目标地址
    struct sockaddr_in dest_addr = {0};
    dest_addr.sin_family = AF_INET;

    in_addr_t iaddr = inet_addr(argv[1]);
    if (iaddr == INADDR_NONE) {
        struct hostent* host = gethostbyname(argv[1]);
        if (NULL == host) {
            printf("无法解析主机名: %s\n", argv[1]);
            close(sfd);
            exit(-1);
        }
        memcpy(&(dest_addr.sin_addr), host->h_addr, host->h_length);
    } else {
        dest_addr.sin_addr.s_addr = iaddr;
    }

    printf("PING %s (%s): %d 数据字节\n",
           argv[1], inet_ntoa(dest_addr.sin_addr), DATA_LEN);

    // 发送和接收数据包
    send_packet(sfd, &dest_addr, sizeof(dest_addr));
    recv_packet(sfd);

    // 打印统计信息
    print_statistics();

    close(sfd);
    return 0;
}

void send_packet(int sfd, struct sockaddr_in* pDestAddr, int len) {
    int packetSize;

    for (nsend = 0; nsend < MAX_PACKETS; nsend++) {
        packetSize = pack(nsend + 1);

        if (sendto(sfd, sendPacket, packetSize, 0,
                  (struct sockaddr*)pDestAddr, len) < 0) {
            printf("第%d个包发送失败: %s\n", nsend + 1, strerror(errno));
            continue;
        }
        printf("第%d个包发送成功\n", nsend + 1);

        sleep(1);
    }
}

int pack(int seq) {
    struct icmp* pIcmp = (struct icmp*)sendPacket;
    memset(pIcmp, 0, sizeof(struct icmp));

    pIcmp->icmp_type = ICMP_ECHO;
    pIcmp->icmp_code = 0;
    pIcmp->icmp_id = getpid();
    pIcmp->icmp_seq = seq;

    // 在数据部分存储发送时间
    struct timeval* tval = (struct timeval*)pIcmp->icmp_data;
    gettimeofday(tval, NULL);

    int packSize = 8 + DATA_LEN;
    pIcmp->icmp_cksum = 0;
    pIcmp->icmp_cksum = ckcsum((u_short*)pIcmp, packSize);

    return packSize;
}

u_short ckcsum(u_short* addr, int len) {
    int nleft = len;
    u_int32_t sum = 0;
    u_short* w = addr;
    u_short answer = 0;

    while (nleft > 1) {
        sum += *w++;
        nleft -= 2;
    }

    if (nleft == 1) {
        *(u_char*)(&answer) = *(u_char*)w;
        sum += answer;
    }

    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);
    answer = ~sum;

    return answer;
}

void recv_packet(int sfd) {
    fd_set readfds;
    struct timeval timeout;
    int max_fd = sfd;
    socklen_t fromLen = sizeof(fromAddr);

    nrecved = 0;

    for (int i = 0; i < MAX_PACKETS; i++) {
        FD_ZERO(&readfds);
        FD_SET(sfd, &readfds);

        timeout.tv_sec = 3;
        timeout.tv_usec = 0;

        int ret = select(max_fd + 1, &readfds, NULL, NULL, &timeout);

        if (ret == 0) {
            printf("请求超时\n");
            continue;
        }

        if (ret < 0) {
            if (errno == EINTR) continue;
            perror("select错误");
            break;
        }

        if (FD_ISSET(sfd, &readfds)) {
            int n = recvfrom(sfd, recvPacket, PACKET_SIZE, 0,
                           (struct sockaddr*)&fromAddr, &fromLen);
            if (n < 0) {
                perror("接收错误");
                continue;
            }

            gettimeofday(&recvTimeVal, NULL);
            if (unpack(n, getpid()) == 0) {
                nrecved++;
            }
        }
    }
}

int unpack(int len, uid_t uid) {
    struct ip* ip_hdr = (struct ip*)recvPacket;
    int iphdrlen = ip_hdr->ip_hl * 4;  // IP头部长度以4字节为单位[6](@ref)

    if (len < iphdrlen + 8) {
        printf("包长度不足\n");
        return -1;
    }

    // 正确计算ICMP头部位置[6](@ref)
    struct icmp* icmp_hdr = (struct icmp*)(recvPacket + iphdrlen);

    if (icmp_hdr->icmp_type == ICMP_ECHOREPLY &&
        icmp_hdr->icmp_id == uid) {

        struct timeval* send_time = (struct timeval*)icmp_hdr->icmp_data;
        struct timeval diff;

        memcpy(&diff, &recvTimeVal, sizeof(diff));
        tv_sub(&diff, send_time);

        double rtt = (diff.tv_sec * 1000.0) + (diff.tv_usec / 1000.0);

        printf("%d 字节来自 %s: icmp_seq=%u ttl=%d 时间=%.1f ms\n",
               len - iphdrlen, inet_ntoa(fromAddr.sin_addr),
               icmp_hdr->icmp_seq, ip_hdr->ip_ttl, rtt);
        return 0;
    }

    return -1;
}

void tv_sub(struct timeval* out, struct timeval* in) {
    if ((out->tv_usec -= in->tv_usec) < 0) {
        --out->tv_sec;
        out->tv_usec += 1000000;
    }
    out->tv_sec -= in->tv_sec;
}

void print_statistics() {
    int lost = nsend - nrecved;
    float loss_rate = (lost * 100.0) / nsend;

    printf("\n--- 统计信息 ---\n");
    printf("%d 包发送, %d 包接收, %.1f%% 丢包\n",
           nsend, nrecved, loss_rate);
}